import itertools
import statsmodels.api as sm
import pandas as pd
from joblib import dump

class Arima():
    def __init__(self):
        self.q = self.d = range(0, 2)
        self.p = range(0, 4)
        self.pdq = list(itertools.product(self.p, self.d, self.q))
        self.seasonal_pdq = [(x[0], x[1], x[2], 12) for x in list(itertools.product(self.p, self.d, self.q))]

    def train(self, train_data):
        AIC = []
        SARIMAX_model = []
        for param in self.pdq:
            for param_seasonal in self.seasonal_pdq:
                try:
                    mod = sm.tsa.statespace.SARIMAX(train_data,
                                                    order=param,
                                                    seasonal_order=param_seasonal,
                                                    enforce_stationarity=False,
                                                    enforce_invertibility=False)

                    results = mod.fit()

                    AIC.append(results.aic)
                    SARIMAX_model.append([param, param_seasonal])
                except:
                    continue

        best_aic_index = AIC.index(min(AIC))
        best_pdq = SARIMAX_model[best_aic_index][0]
        best_seasonal_pdq = SARIMAX_model[best_aic_index][1]

        print(f'The smallest AIC is {min(AIC)} for model SARIMAX{best_pdq}x{best_seasonal_pdq}')

        best_mod = sm.tsa.statespace.SARIMAX(train_data,
                                             order=best_pdq,
                                             seasonal_order=best_seasonal_pdq,
                                             enforce_stationarity=False,
                                             enforce_invertibility=False)

        best_results = best_mod.fit()

        dump(best_results, 'arima_model.joblib')

        return best_results
